import sqlite3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def sqlConn(dbFile):
    conn = sqlite3.connect(dbFile)
    return conn


def computeKDJ(dbFile):
    # with sqlite3.connect('./sqlData.db') as  conn:
    conn = sqlConn(dbFile=dbFile)
    data = []
    cur = conn.cursor()
    cur.execute("select * from jzfg where date>'2016-08-09'")
    fieldList = [t[0] for t in cur.description]# get fields from define table
    while True:
        row = cur.fetchone()
        if type(row) == type(None):
            break
        data.append(row)
    index = np.array(data)[:,1]
    dataList = pd.DataFrame(data,index=index,columns=fieldList,dtype=float)
    lowList = dataList['low'].rolling(window=9).min()
    highList = dataList['high'].rolling(window=9).max()
    rsv = (dataList['close'] - lowList)/(highList-lowList)*100
    dfData = pd.DataFrame()
    dfData['high'] = dataList['high']
    dfData['low'] = dataList['low']
    dfData['close'] = dataList['close']
    dfData['K'] = rsv.ewm(com=2).mean()
    dfData['D'] = dfData['K'].ewm(com=2).mean()
    dfData['J'] = 3*dfData['K'] - 2*dfData['D']
    dfData.index = dataList['date'].values
    dfData.index.name = 'date'
    dfData = dfData.dropna()
    dfData.to_excel('D:/hello.xlsx')
    dfData[['K','D','J']].plot(title='KDJ')
    plt.show()


def computeMACD():

    pass

if __name__ == '__main__':
    computeKDJ('./sqlData.db')